import asyncio
import json
from pprint import pprint

from agents.requirement_analyzer import agent_requirements, RequirementDSL
from agents.use_case_generator import agent_use_case, UseCaseDSL
from agents.ssd_generator import agent_ssd, SSDSL
from agents.class_diagram_generator import agent_class_diagram, ClassDiagramDSL
from agents.ocl_generator import agent_ocl, OCDSL

from openai import AsyncOpenAI
from openai import Agent, Runner, RunConfig, OpenAIProvider


async def main():
    client = AsyncOpenAI(
        api_key="YOUR_OPENAI_API_KEY",  # 替换为你的API Key
    )

    provider = OpenAIProvider(openai_client=client)

    user_input = input("请输入图书馆管理系统的需求描述：")

    # Step 1: 需求分析
    req_result = await Runner.run(agent_requirements, input=user_input, model_provider=provider)
    req_dsl = req_result.final_output_as(RequirementDSL)

    print("\n【需求分析结果】")
    pprint(req_dsl)

    # Step 2: 用例图生成
    uc_result = await Runner.run(agent_use_case, input=req_dsl.summary, model_provider=provider)
    uc_dsl = uc_result.final_output_as(UseCaseDSL)

    print("\n【用例图DSL】")
    pprint(uc_dsl)

    # Step 3: 系统顺序图生成
    ssd_result = await Runner.run(agent_ssd, input=req_dsl.summary, model_provider=provider)
    ssd_dsl = ssd_result.final_output_as(SSDSL)

    print("\n【系统顺序图DSL】")
    pprint(ssd_dsl)

    # Step 4: 类图生成
    cd_result = await Runner.run(agent_class_diagram, input=req_dsl.summary, model_provider=provider)
    cd_dsl = cd_result.final_output_as(ClassDiagramDSL)

    print("\n【概念类图DSL】")
    pprint(cd_dsl)

    # Step 5: OCL 合约生成
    ocl_result = await Runner.run(agent_ocl, input=req_dsl.summary, model_provider=provider)
    ocl_dsl = ocl_result.final_output_as(OCDSL)

    print("\n【OCL 合约DSL】")
    pprint(ocl_dsl)

    # Step 6: 整合所有DSL并输出
    full_model = {
        "requirement": req_dsl.__dict__,
        "use_cases": uc_dsl.__dict__,
        "sequence_diagrams": ssd_dsl.__dict__,
        "class_diagrams": cd_dsl.__dict__,
        "ocl_contracts": ocl_dsl.__dict__
    }

    with open("library_system_model.json", "w", encoding="utf-8") as f:
        json.dump(full_model, f, ensure_ascii=False, indent=2)

    print("\n✅ 最终整合模型已生成！")
    print("已保存至 library_system_model.json")


if __name__ == "__main__":
    asyncio.run(main())